package linkage;

import hmm.MarkovianProbabilities;
import static java.lang.Math.pow;


public class Probabilities implements MarkovianProbabilities {
	
	public Probabilities(int[][] pedigree, byte[][] loci, int numOfNonFounders, double[] markersDist) {
		this.pedigree = pedigree;
		this.loci = loci;
		this.numOfNonFounders = numOfNonFounders;
		this.markersDist = markersDist;
		nonFounders = new int[numOfNonFounders];
		int id = 0;
		for(int i = 0; i < numOfNonFounders; i++) {
			while(pedigree[id][fatherCol] == noParent && pedigree[id][motherCol] == noParent) { id++; }
			nonFounders[i] = id;
		}
	}

	@Override
	public double transitionProb(int k, int l, int i) {
		/*
		 * in this case, return the probability of recombination and no recombination
		 * between marker i and marker i+1 according to the inheritance vector
		 */
		
		if(i<0) { return 1; } // in case this is the first marker
		
		// how much recombinations ware between state k and l
		int recombinationsCount = numberOfSetBits(k ^ l);
		return pow(markersDist[i], recombinationsCount) * 
				pow(1 - markersDist[i], numOfNonFounders - recombinationsCount);
	}

	@Override
	public double observationProb(int state, int observationIndex) {
		/*
		 *  ASSUMPTION: there are only two possible alleles, and the
		 *  probability of each is 0.5
		 */
		double[] p = new double[pedigree.length];
		int id, father, mother;
		for(int i = 0; i < numOfNonFounders; i++) {
			id = nonFounders[i];
			father = pedigree[id][fatherCol];
			mother = pedigree[id][motherCol];
			if(father == -1 && mother == -1) {
				p[id] = inheritanceProbability(id, observationIndex);
			} else if (father >= 0 && mother >= 0) {
				p[id] = inheritanceProbability(id, father, mother, observationIndex);
			}
			if(father != -1 && mother != -1) {
				p[id] = p[id] * p[father] * p[mother];
			}
		}
		double sum = 0;
		for(int j = 0; j < pedigree.length; j++) {
			sum += p[j];
		}
		return sum;
	}
	
	/*
	 * ================= private methods
	 */
	
	/*
	 * return the probability that the child inherited the marker from his parents
	 * this function is used for non-founders.
	 */
	private double inheritanceProbability(int id, int father, int mother, int marker) {
		double p = 0;
		int allele1 = marker*2, allele2 = marker*2+1;
		
		// assume that the first locus is paternal and the second is maternal
		if((loci[id][allele1] == loci[father][allele1] || loci[id][allele1] == loci[father][allele2])
				&& (loci[id][allele2] == loci[mother][allele1] || loci[id][allele2] == loci[mother][allele2])) {
			double tmp = 0;
			if(loci[father][allele1] == loci[father][allele2]) { tmp = 0.5; } // 0.5 * 1
			else { tmp = 0.25; } // 0.5 * 0.5
			if(loci[mother][allele1] == loci[mother][allele2]) { tmp *= 0.5; } // tmp * 0.5 * 1
			else { tmp *= 0.25; } // tmp * 0.5 * 0.5
			p += tmp;
		}
		// assume that the first locus is maternal and the second is paternal
		if((loci[id][allele2] == loci[father][allele1] || loci[id][allele2] == loci[father][allele2])
				&& (loci[id][allele1] == loci[mother][allele1] || loci[id][allele1] == loci[mother][allele2])) {
			double tmp = 0;
			if(loci[father][allele1] == loci[father][allele2]) { tmp = 0.5; } // 0.5 * 1
			else { tmp = 0.25; } // 0.5 * 0.5
			if(loci[mother][allele1] == loci[mother][allele2]) { tmp *= 0.5; } // tmp * 0.5 * 1
			else { tmp *= 0.25; } // tmp * 0.5 * 0.5
			p += tmp;
		}
		
		return p;
	}
	
	/*
	 * return the probability that the child inherited the marker from his parents
	 * this function is used for founders only.
	 */
	private double inheritanceProbability(int id, int marker) {
		if(loci[id][marker*2] == loci[id][marker*2+1]) {
			return 0.25; // 0.5 * 0.5
		} else {
			return 0.5; // 0.5 * 0.5 + 0.5 * 0.5
		}
	}
	
	/*
	 * calculate the number of "1" bits in a number
	 */
	private int numberOfSetBits(int num) {
		num = num - ((num >> 1) & 0x55555555);
		num = (num & 0x33333333) + ((num >> 2) & 0x33333333);
	    return (((num + (num >> 4)) & 0x0F0F0F0F) * 0x01010101) >> 24;
	}


	/*
	 * ================= private variables
	 */
	// the data of the pedigree. the structure of the matrix is:
	// [id], [father, mother, sex, affection status].
	// for unknown parent - use (-1); sex - m = 0, female = 1;
	// affection status - 0 = unknown, 1 = unaffected; 2 - affected.
	// ASSUMPTION: ID(parent) < ID(child)
	private int[][] pedigree;
	// the data of the loci, the structure is [id][locus]. alleles start with 1, unknown alleles are 0.
	private byte[][] loci;
	// array of distances of markers. the cell i indicates the distance
	// between the i'th and the (i+1)'th markers.
	private double[] markersDist = {0.3};
	private int numOfNonFounders; // number of non-founders in the pedigree
	private int[] nonFounders; // array of all the non-founders.
	
	/*
	 * ================= constants
	 */
	private final int fatherCol = 0;
	private final int motherCol = 1;
	private final int sexCol = 2;
	private final int affectionStatusCol = 3;
	
	/**
	 * if the value noParent is in the father / mother column in the pedigree - this means that the parent doesn't exists in the pedigree
	 */
	public final int noParent = -1;
}
